import os
import torch
from application.config.config import Config
from application.model.rnn_model import RNN
from .train_valid import TrainValid

class Predict:
    def __init__(self):
        self.config = Config()
        self.train_valid = TrainValid()
        self.rnn = RNN()
        self.hidden = self.rnn.init_hidden()

    def _test(self, text_tensor):
        """
        计算预测结果函数
        :param text_tensor:
        :return:
        """
        self.rnn.load_state_dict(torch.load(self.config.MODEL_SAVE_PATH))
        with torch.no_grad():
            for i in range(text_tensor.size()[1]):
                output, hidden = self.rnn(text_tensor[0][i].unsqueeze(0), self.hidden)
        return output

    def model_predict(self, input_text):
        """
        把输入的症状的文本进行计算得到预测标签
        :param input_text: 症状的文本
        :return: 标签
        """
        with torch.no_grad():
            output = self._test(self.train_valid.get_bert_encode_for_single(input_text))
            _, topi = output.topk(1, 1)  # topk 第一个参数是k的个数，是哪个维度
            return topi.item()

    def batch_predict(self):
        """
        批量审核每个疾病的症状
        :param input_path: 待审核疾病文件所在的目录
        :param output_path: 审核完成后疾病文件所在的目录
        :return:
        """
        csv_list = os.listdir(self.config.NOREVIEW_PATH)
        for disease in csv_list:
            with open(os.path.join(self.config.NOREVIEW_PATH, disease), 'r', encoding="utf-8") as fr:
                with open(os.path.join(self.config.REVIEWED_PATH, disease), 'w') as fw:
                    lines = fr.readlines()
                    for line in lines:
                        result = self.model_predict(line)
                        if result:
                            fw.write(line)
                        else:
                            pass
